/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2020-2023. All rights reserved.
 * Description: hccl实现基类文件
 */

#include <atomic>
#include <chrono>
#include <thread>
#include <algorithm>
#include <numeric>
#include <unordered_set>
#include <sys/time.h>
#include "slog.h"
#include "slog_api.h"
#include "dlprof_function.h"
#include "dlrt_function.h"
#include "dltdt_function.h"
#include "externalinput_pub.h"
#include "p2p_mgmt_pub.h"
#include "../command_handle.h"
#include "config.h"
#include "transport_heterog_event_roce_pub.h"
#include "transport_heterog_event_tcp_pub.h"
#include "transport_heterog_roce_pub.h"
#include "transport_roce.h"
#include "alg_remote_lookup.h"
#include "alg_coll_remote_update.h"
#include "alg_coll_remote_lookup.h"
#include "callback_thread_manager.h"
#include "transport_shm_event_pub.h"
#include "stream_active_manager.h"
#include "device_capacity.h"
#include "profiling_manager.h"
#include "task_exception_handler_pub.h"
#include "hccl_impl_cpu.h"

using namespace std;

namespace hccl {
HcclImplCpu::HcclImplCpu():HcclImplBase()
{};

HcclImplCpu::~HcclImplCpu()
{};


HcclResult HcclImplCpu::Send(const std::string &tag, void *inputPtr, u64 count, HcclDataType dataType,
    u32 destRank, rtStream_t stream)
{
    if (!IsAtomicInit()) {
        HCCL_ERROR("[HcclImplBase][Send]errNo[0x%016llx] hccl init must be called before call this function",
            HCCL_ERROR_CODE(HCCL_E_UNAVAIL));
        return HCCL_E_UNAVAIL;
    }

    Stream streamObj(stream);
    CHK_RET(callbackTask_->CallbackRegStream(stream));

    CHK_RET(implAlg_->Send(tag, inputPtr, count, dataType, destRank, streamObj));

    return HCCL_SUCCESS;
}

HcclResult HcclImplCpu::Receive(const std::string &tag, void *outputPtr, u64 count, HcclDataType dataType,
    u32 srcRank, rtStream_t stream)
{
    if (!IsAtomicInit()) {
        HCCL_ERROR("[HcclImplBase][Receive]errNo[0x%016llx] hccl init must be called before call this function",
            HCCL_ERROR_CODE(HCCL_E_UNAVAIL));
        return HCCL_E_UNAVAIL;
    }

    Stream streamObj(stream);
    CHK_RET(callbackTask_->CallbackRegStream(stream));
    CHK_RET(implAlg_->Receive(tag, outputPtr, count, dataType, srcRank, streamObj));
    return HCCL_SUCCESS;
}


HcclResult HcclImplCpu::BcastByOS(void* buf, uint64_t count, HcclDataType dataType,
    int root, uint64_t flag, const std::string &tag)
{
    if (!IsAtomicInit()) {
        HCCL_ERROR(
            "[HcclImplBase][Broadcast]errNo[0x%016llx] hccl init must be called before call this function",
            HCCL_ERROR_CODE(HCCL_E_UNAVAIL));
        return HCCL_E_UNAVAIL;
    }

    /* ------------集合通信资源准备------------ */
    if (opTypeTagMap_.find(HCOM_OP_TYPE_BROADCAST) == opTypeTagMap_.end()) {
        HCCL_ERROR("[HcclImplBase][BcastByOS]tag is not found");
        return HCCL_E_PARA;
    }

    if (LIKELY(!(implAlg_->IsExistCommRes(tag)))) {
        HCCL_ERROR("[HcclImplBase][BcastByOS]comm is not exist, please call 'HcomPrepareStart'");
        return HCCL_E_PARA;
    }

    /*  ------------执行算法-------------- */
    Stream streamObj(nullptr);
    HcclResult ret =
        implAlg_->Broadcast(tag, buf, count, dataType, root, streamObj);
    CHK_PRT_RET(ret == HCCL_E_AGAIN,
        HCCL_WARNING("[HcclImplBase][Broadcast]group has been destroyed. Break!"), ret);
    CHK_PRT_RET(ret != HCCL_SUCCESS,
        HCCL_ERROR("[HcclImplBase][Broadcast]errNo[0x%016llx] tag[%s],braodcast op run failed",
        HCCL_ERROR_CODE(ret),
            tag.c_str()), ret);
    return HCCL_SUCCESS;
}

HcclResult HcclImplCpu::GatherByOs(void* inputBuf, uint64_t inputCount, HcclDataType inputType, void* outputBuf,
    uint64_t outputCount, HcclDataType outputType, int root, uint64_t flag, const std::string &tag)
{
    if (!IsAtomicInit()) {
        HCCL_ERROR("[HcclImplBase][GatherByOs]errNo[0x%016llx] hccl init must be called before call this" \
            " function",
            HCCL_ERROR_CODE(HCCL_E_UNAVAIL));
        return HCCL_E_UNAVAIL;
    }

    CHK_PRT_RET(inputType != outputType, HCCL_ERROR("[HcclImplBase][GatherByOs]input data type[%s] is not" \
        " same with output data type[%s]", GetDataTypeEnumStr(inputType).c_str(),
        GetDataTypeEnumStr(outputType).c_str()), HCCL_E_PARA);

    /* ------------集合通信资源准备------------ */

    if (opTypeTagMap_.find(HCOM_OP_TYPE_GATHER) == opTypeTagMap_.end()) {
        HCCL_ERROR("[HcclImplBase][GatherByOs]tag is not found");
        return HCCL_E_PARA;
    }

    if (LIKELY(!(implAlg_->IsExistCommRes(tag)))) {
        HCCL_ERROR("[HcclImplBase][GatherByOs]comm is not exist, please call 'HcomPrepareStart'");
        return HCCL_E_PARA;
    }

    /*  ------------执行算法-------------- */
    auto startut = TIME_NOW();
    Stream streamObj(nullptr);
    HcclResult ret = implAlg_->Gather(tag, inputBuf, outputBuf, root, inputCount, inputType, streamObj);
    CHK_PRT_RET(ret != HCCL_SUCCESS,
        HCCL_ERROR("[HcclImplBase][GatherByOs]errNo[0x%016llx] tag[%s],gather run failed",
            HCCL_ERROR_CODE(ret), tag.c_str()), ret);
    HCCL_INFO("tag[%s],GatherByOs run success,take time [%lld]us", tag.c_str(), DURATION_US(TIME_NOW() - startut));
    return HCCL_SUCCESS;
}

HcclResult HcclImplCpu::PrepareStart(const std::string &tag, void *inputPtr, u64 count, HcclDataType dataType,
    u32 peerRank)
{
    CHK_RET(notifyPool_->RegisterOp(tag));
    if (!prepareCommThread_) {
        prepareCommThread_.reset(new (std::nothrow) std::thread(&HcclImplBase::PrepareCommAsync, this));
        CHK_PRT_RET(!prepareCommThread_, HCCL_ERROR("[PrepareStart]create threads failed "), HCCL_E_INTERNAL);
    }
    HcclOpCommDesc opDesc(inputPtr, count, dataType, peerRank);
    std::unique_lock<std::mutex> requestLock(prepareRequestMutex_);
    prepareRequest_.insert({tag, opDesc});
    requestLock.unlock();
    NotifyPrepareComm();
    return HCCL_SUCCESS;
}

HcclResult HcclImplCpu::PrepareStartColl(
    const std::string &tag, const HcomCollOpInfo &opInfo, HcomOperationType opType)
{
    CHK_RET(DlTdtFunction::GetInstance().DlTdtFunctionHeterogInit());
    CHK_RET(notifyPool_->RegisterOp(tag));
    opTypeTagMap_[opType] = tag;
    if (opType == HCOM_OP_TYPE_GATHER) {
        CHK_RET(PrepareGather(tag, opInfo, opType));
    } else if (opType == HCOM_OP_TYPE_BROADCAST) {
        CHK_RET(PrepareBroadcast(tag, opInfo, opType));
    }
    CHK_RET(notifyPool_->UnregisterOp(tag));
    return HCCL_SUCCESS;
}

HcclResult HcclImplCpu::PrepareQuerry(const std::string &tag, u32& status)
{
    CHK_PRT_RET(!prepareCommThread_,
        HCCL_ERROR("[PrepareQuerry]PrepareCommAsync threads not exist."), HCCL_E_INTERNAL);

    std::unique_lock<std::mutex> requestLock(prepareRequestMutex_);
    auto iter = prepareRequest_.find(tag);
    if (iter == prepareRequest_.end()) {
        HCCL_ERROR("opTag[%s] is not exsit.", tag.c_str());
        return HCCL_E_PARA;
    } else {
        status = (iter->second.status == 0xffffffff) ? 1 : iter->second.status;
    }
    requestLock.unlock();
    if (status == HETEROG_P2P_SUCCESS) {
        CHK_RET(notifyPool_->UnregisterOp(tag));
    }
    return HCCL_SUCCESS;
}

HcclResult HcclImplCpu::PrepareCancel(const std::string &tag, u32& status)
{
    CHK_PRT_RET(!prepareCommThread_,
        HCCL_ERROR("[PrepareCancel]PrepareCommAsync threads not exist."), HCCL_E_INTERNAL);

    std::unique_lock<std::mutex> requestLock(prepareRequestMutex_);
    requestLock.lock();
    auto iter = prepareRequest_.find(tag);
    if (iter == prepareRequest_.end()) {
        HCCL_WARNING("prepareRequest_ opTag[%s] is not exsit.", tag.c_str());
    } else {
        status = iter->second.status;
        prepareRequest_.erase(iter);
    }
    requestLock.unlock();

    if (implAlg_ != nullptr) {
        implAlg_->CancelCommRes(tag);
    }

    return HCCL_SUCCESS;
}

HcclResult HcclImplCpu::PrepareGather(const std::string &tag, const HcomCollOpInfo &opInfo,
    HcomOperationType opType)
{
    CHK_SMART_PTR_NULL(implAlg_);
    u32 perDataSize = 0;
    CHK_RET(SalGetDataTypeSize(opInfo.dataType, perDataSize));

    DeviceMem inputMem = DeviceMem::create(opInfo.inputAddr, opInfo.count * perDataSize);
    DeviceMem outputMem;

    if (opInfo.root != userRank_) {
        if (!isServerInter_) {
            u64 workSpaceMemSize = opInfo.count * perDataSize * userRankSize_;
            if (tagWorkSpaceMem_.find(tag) == tagWorkSpaceMem_.end()) {
                tagWorkSpaceMem_[tag] = HostMem::alloc(workSpaceMemSize, false);
                CHK_SMART_PTR_NULL(tagWorkSpaceMem_[tag].ptr());
            }
            void *workSpaceMem = tagWorkSpaceMem_[tag].ptr();
            CHK_PTR_NULL(workSpaceMem);
            outputMem = DeviceMem::create(workSpaceMem, workSpaceMemSize);
        } else {
            auto outputMemPtr = opInfo.outputAddr == nullptr ? opInfo.inputAddr : opInfo.outputAddr;
            outputMem = DeviceMem::create(outputMemPtr, opInfo.count * perDataSize);
        }
    } else {
        outputMem = DeviceMem::create(opInfo.outputAddr, userRankSize_ * opInfo.count * perDataSize);
    }
    HCCL_DEBUG("inputMem[%p],outputMem[%p].", inputMem.ptr(), outputMem.ptr());

    CHK_RET(IsHostUseDevNic(isHostUseDevNic_));
    if (isHaveCpuRank_ && !isSetHDCModeInfo_ && isServerInter_) {
        if (static_cast<s32>(devicePhyId_) == HOST_DEVICE_ID && isHostUseDevNic_) {
            CHK_RET(SetDevice(true));
            CHK_RET(InitPsHdcRaResource());
            CHK_RET(InitHdcPsRdmaHandle());
            isSetHDCModeInfo_ = true;
        } else {
            CHK_RET(InitPsRaResource(deviceLogicId_, devicePhyId_, hostIp_, ranksPort_[realUserRank_], false));
            CHK_RET(InitPsRdmaHandle(deviceLogicId_, devicePhyId_, hostIp_));
        }
    }
    CHK_RET(MrManagerInit());
    CHK_RET(InitRecvMsgAndRequestBuffer());
    CHK_RET(InitMemBlocksAndRecvWrMem());
    implAlg_->SetHDCModeInfo(rankDevicePhyIdNicInfoMap_, groupRanksPort_, isSetHDCModeInfo_, isUseRankPort_);

    // 创建通信域
    AlgType algType;
    algType.algoLevel0 = AlgTypeLevel0::ALG_LEVEL0_NP_STAR;
    algType.algoLevel1 = AlgTypeLevel1::ALG_LEVEL1_STAR;
    CHK_RET(implAlg_->SetAlgType(algType, HcclCMDType::HCCL_CMD_GATHER));
    HcclResult ret = implAlg_->CreateComm(tag, inputMem, outputMem, algType, opInfo.root);
    CHK_PRT_RET(ret != HCCL_SUCCESS,
        HCCL_ERROR("[HcclImplBase][AllGather]errNo[0x%016llx] tag[%s],all gather create comm failed",
        HCCL_ERROR_CODE(ret), tag.c_str()), ret);

    return HCCL_SUCCESS;
}

HcclResult HcclImplCpu::PrepareBroadcast(
    const std::string &tag, const HcomCollOpInfo &opInfo, HcomOperationType opType)
{
    CHK_SMART_PTR_NULL(implAlg_);
    u32 perDataSize = 0;
    CHK_RET(SalGetDataTypeSize(opInfo.dataType, perDataSize));
    DeviceMem devMem = DeviceMem::create(opInfo.inputAddr, opInfo.count * perDataSize);

    CHK_RET(IsHostUseDevNic(isHostUseDevNic_));
    if (isHaveCpuRank_ && !isSetHDCModeInfo_ && isServerInter_) {
        if (static_cast<s32>(devicePhyId_) == HOST_DEVICE_ID && isHostUseDevNic_) {
            CHK_RET(SetDevice(true));
            CHK_RET(InitPsHdcRaResource());
            CHK_RET(InitHdcPsRdmaHandle());
            isSetHDCModeInfo_ = true;
        } else {
            CHK_RET(InitPsRaResource(deviceLogicId_, devicePhyId_, hostIp_, ranksPort_[realUserRank_], false));
            CHK_RET(InitPsRdmaHandle(deviceLogicId_, devicePhyId_, hostIp_));
        }
    }
    CHK_RET(MrManagerInit());
    CHK_RET(InitRecvMsgAndRequestBuffer());
    CHK_RET(InitMemBlocksAndRecvWrMem());

    implAlg_->SetHDCModeInfo(rankDevicePhyIdNicInfoMap_, groupRanksPort_, isSetHDCModeInfo_, isUseRankPort_);
    AlgType algType;
    algType.algoLevel0 = AlgTypeLevel0::ALG_LEVEL0_NP_STAR;
    algType.algoLevel1 = AlgTypeLevel1::ALG_LEVEL1_STAR;
    CHK_RET(implAlg_->SetAlgType(algType, HcclCMDType::HCCL_CMD_BROADCAST));
    CHK_RET(implAlg_->CreateComm(tag, devMem, devMem, algType, opInfo.root));

    return HCCL_SUCCESS;
}

HcclResult HcclImplCpu::GetHeterogCollTag(HcomOperationType opType, std::string &tag)
{
    if (opTypeTagMap_.find(opType) == opTypeTagMap_.end()) {
        return HCCL_E_NOT_FOUND;
    } else {
        tag = opTypeTagMap_[opType];
    }
    return HCCL_SUCCESS;
}

}